
import cv2
import numpy as np

def get_correspondences_using_optical_flow(img1,img2):

    # Lucas-Kanade光流算法的参数
    lk_params = dict(winSize=(15, 15),  # 窗口大小
                     maxLevel=2,
                     criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))

    ##获得img1的特征点
    fast = cv2.FastFeatureDetector_create()
    kps1 = fast.detect(img1, None)  # 返回的是keyPoint
    # 取其中的pt
    pts1 = []
    for kp in kps1:
        pts1.append(kp.pt)
    pts1 = np.array(pts1, dtype=np.float32)
    pts1 = pts1.ravel()
    pts1 = pts1.reshape(-1, 1, 2)

    ##计算光流，
    ##p1和点数量和p0一样
    ##st是和p1,p0数量一样的向量，1表示p0对应的p1找到了
    ##err返回匹配误差
    pts2,st,err=cv2.calcOpticalFlowPyrLK(img1,img2,pts1,None,**lk_params)

    #选择好的匹配点
    #good_pts1=pts1[(st==1)&(err<=10)]  #取好的匹配  得到的形式：[ [x1,y1],[x2,y2], ... ,[xn,yn]  ]
    #good_pts2=pts2[(st==1)&(err<=10)]
    good_pts1=pts1[st==1]  #取好的匹配  得到的形式：[ [x1,y1],[x2,y2], ... ,[xn,yn]  ]
    good_pts2=pts2[st==1]

    return good_pts1,good_pts2


#####test
img1 = cv2.imread('../misc_pic/yl01_01.jpg')
img2 = cv2.imread('../misc_pic/yl01_02.jpg')
gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

pts1,pts2=get_correspondences_using_optical_flow(gray1,gray2)
N=len(pts1)

#创建一些随机颜色
color=np.random.randint(0,255,(N,3))
# 绘制光流轨迹
mask=np.zeros_like(img1)
for i, (pt1, pt2) in enumerate(zip(pts1, pts2)):
    a, b = pt1.ravel()  # 阵列平坦化
    c, d = pt2.ravel()
    # 绘制光流方向线段
    mask = cv2.line(mask, (a, b), (c, d), color[i].tolist(), 1)
    frame = cv2.circle(img1, (a, b), 2, color[i].tolist(), -1)
img = cv2.add(img1, mask)
while True:
    cv2.imshow('frame', img)
    k = cv2.waitKey(30) & 0xff
    if k == 27:
        break
